#include <wdm.h>

#include "driver.h"
#include "device.h"
#include "log.h"
#include "pool.h"

using namespace windows::kernal;

//#define DEVICE_NAME L"\\Device\\Asynchronous"
//#define SYMBOLIC_LINK_NAME L"\\DosDevices\\Asynchronous"

#define DEVICE_NAME L"Asynchronous"
#define SYMBOLIC_LINK_NAME L"Asynchronous"

Device *device = NULL;

NTSTATUS DriverEntry(
  _In_ PDRIVER_OBJECT driver_object, 
  _In_ PUNICODE_STRING registry_path) {

  ENTER();

  NTSTATUS status = STATUS_UNSUCCESSFUL;
  UNREFERENCED_PARAMETER(registry_path);
  driver_object->DriverUnload = DriverUnload;
  driver_object->MajorFunction[IRP_MJ_CREATE] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_CREATE_NAMED_PIPE] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_CLOSE] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_READ] = DriverReadDispatch;
  driver_object->MajorFunction[IRP_MJ_WRITE] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_QUERY_INFORMATION] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_SET_INFORMATION] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_QUERY_EA] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_SET_EA] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_FLUSH_BUFFERS] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_QUERY_VOLUME_INFORMATION] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_SET_VOLUME_INFORMATION] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_DIRECTORY_CONTROL] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_FILE_SYSTEM_CONTROL] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_INTERNAL_DEVICE_CONTROL] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_SHUTDOWN] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_LOCK_CONTROL] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_CLEANUP] = DriverCleanupDispatch;
  driver_object->MajorFunction[IRP_MJ_CREATE_MAILSLOT] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_QUERY_SECURITY] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_SET_SECURITY] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_POWER] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_SYSTEM_CONTROL] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_DEVICE_CHANGE] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_QUERY_QUOTA] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_SET_QUOTA] = DriverDefaultDispatch;
  driver_object->MajorFunction[IRP_MJ_PNP] = DriverPnpDispath;

  device = new Device(driver_object);
  if (device) {
    status = device->Create(DEVICE_NAME, SYMBOLIC_LINK_NAME);
    if (!NT_SUCCESS(status)) {
      ERR("DriverEntry device->Initialize fail, status=%#X\n", status);
    }
  } else {
    status = STATUS_INSUFFICIENT_RESOURCES;
    ERR("DriverEntry new Device fail\n");
  }

  LEAVE();
  return status;
}

VOID DriverUnload(_In_ PDRIVER_OBJECT driver_object) {
  ENTER();

  UNREFERENCED_PARAMETER(driver_object);

  if (device) {
    delete device;
    device = NULL;
  }
  LEAVE();
}

NTSTATUS DriverDefaultDispatch(_In_ PDEVICE_OBJECT device_object, _In_ PIRP irp) {
  PAGED_CODE();
  ENTER();

  UNREFERENCED_PARAMETER(device_object);
  static LPCSTR irp_major_type[] = {
    "IRP_MJ_CREATE",
    "IRP_MJ_CREATE_NAMED_PIPE",
    "IRP_MJ_CLOSE",
    "IRP_MJ_READ",
    "IRP_MJ_WRITE",
    "IRP_MJ_QUERY_INFORMATION",
    "IRP_MJ_SET_INFORMATION",
    "IRP_MJ_QUERY_EA",
    "IRP_MJ_SET_EA",
    "IRP_MJ_FLUSH_BUFFERS",
    "IRP_MJ_QUERY_VOLUME_INFORMATION",
    "IRP_MJ_SET_VOLUME_INFORMATION",
    "IRP_MJ_DIRECTORY_CONTROL",
    "IRP_MJ_FILE_SYSTEM_CONTROL",
    "IRP_MJ_DEVICE_CONTROL",
    "IRP_MJ_INTERNAL_DEVICE_CONTROL",
    "IRP_MJ_SHUTDOWN",
    "IRP_MJ_LOCK_CONTROL",
    "IRP_MJ_CLEANUP",
    "IRP_MJ_CREATE_MAILSLOT",
    "IRP_MJ_QUERY_SECURITY",
    "IRP_MJ_SET_SECURITY",
    "IRP_MJ_POWER",
    "IRP_MJ_SYSTEM_CONTROL",
    "IRP_MJ_DEVICE_CHANGE",
    "IRP_MJ_QUERY_QUOTA",
    "IRP_MJ_SET_QUOTA",
    "IRP_MJ_PNP",
  };

  PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(irp);
  if (_ARRAYSIZE(irp_major_type) <= stack->MajorFunction) {
    ERR("DriverDefaultDisplay major(%d) is unexpect\n", irp_major_type[stack->MajorFunction]);
  } else {
    INFO("DriverDefaultDisplay major(%d)\n", stack->MajorFunction);
  }
  irp->IoStatus.Status = STATUS_SUCCESS;
  irp->IoStatus.Information = 0;
  IoCompleteRequest(irp, IO_NO_INCREMENT);
  
  LEAVE();
  return STATUS_SUCCESS;
}

NTSTATUS DriverPnpDispath(_In_ PDEVICE_OBJECT device_object, _In_ PIRP irp) {
  PAGED_CODE();
  ENTER();

  UNREFERENCED_PARAMETER(device_object);
  static LPCSTR pnp_type[] = {
    "IRP_MN_START_DEVICE",
    "IRP_MN_QUERY_REMOVE_DEVICE",
    "IRP_MN_REMOVE_DEVICE",
    "IRP_MN_CANCEL_REMOVE_DEVICE",
    "IRP_MN_STOP_DEVICE",
    "IRP_MN_QUERY_STOP_DEVICE",
    "IRP_MN_CANCEL_STOP_DEVICE",
    "IRP_MN_QUERY_DEVICE_RELATIONS",
    "IRP_MN_QUERY_INTERFACE",
    "IRP_MN_QUERY_CAPABILITIES",
    "IRP_MN_QUERY_RESOURCES",
    "IRP_MN_QUERY_RESOURCE_REQUIREMENTS",
    "IRP_MN_QUERY_DEVICE_TEXT",
    "IRP_MN_FILTER_RESOURCE_REQUIREMENTS",
    ""
    "IRP_MN_READ_CONFIG",
    "IRP_MN_WRITE_CONFIG",
    "IRP_MN_EJECT",
    "IRP_MN_SET_LOCK",
    "IRP_MN_QUERY_ID",
    "IRP_MN_QUERY_PNP_DEVICE_STATE",
    "IRP_MN_QUERY_BUS_INFORMATION",
    "IRP_MN_DEVICE_USAGE_NOTIFICATION",
    "IRP_MN_SURPRISE_REMOVAL",
    "",
    #if (NTDDI_VERSION >= NTDDI_WIN7)
    "IRP_MN_DEVICE_ENUMERATED",
    #endif
  };
  
  PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(irp);
  if (_ARRAYSIZE(pnp_type) <= stack->MinorFunction) {
    ERR("DriverPnpDispath minor(%s) is unexpect\n", pnp_type[stack->MajorFunction]);
  } else {
    INFO("DriverPnpDispath minor(%d)\n", stack->MajorFunction);
  }
  irp->IoStatus.Status = STATUS_SUCCESS;
  irp->IoStatus.Information = 0;
  IoCompleteRequest(irp, IO_NO_INCREMENT);

  LEAVE();
  return STATUS_SUCCESS;
}

NTSTATUS DriverReadDispatch(_In_ PDEVICE_OBJECT device_object, _In_ PIRP irp) {
  PAGED_CODE();
  ENTER();
  UNREFERENCED_PARAMETER(device_object);

  INFO("DriverReadDispatch device_object=%p, irp=%p\n", device_object, irp);
  IoSetCancelRoutine(irp, Device::CancelRoutine);
  IoMarkIrpPending(irp);
  LEAVE();
  return STATUS_PENDING;
}

NTSTATUS DriverCleanupDispatch(_In_ PDEVICE_OBJECT device_object, _In_ PIRP irp) {
  NTSTATUS status;

  PAGED_CODE();
  ENTER();
  UNREFERENCED_PARAMETER(device_object);

  status = device->CompleteAllPendingIrp();
  if (!NT_SUCCESS(status)) {
    ERR("DriverCleanupDispatch CompleteAllPendingIrp fail, status=%#X\n", status);
  }
  irp->IoStatus.Status = STATUS_SUCCESS;
  irp->IoStatus.Information = 0;
  IoCompleteRequest(irp, IO_NO_INCREMENT);

  LEAVE();
  return STATUS_SUCCESS;
}